package com.fastjrun.codeg.javamapper;

import com.fastjrun.codeg.javamapper.elements.sqlprovider.ProviderInsertBatchMethodGenerator;
import org.mybatis.generator.api.dom.java.CompilationUnit;
import org.mybatis.generator.api.dom.java.TopLevelClass;
import org.mybatis.generator.codegen.mybatis3.javamapper.SqlProviderGenerator;
import org.mybatis.generator.codegen.mybatis3.javamapper.elements.sqlprovider.AbstractJavaProviderMethodGenerator;

import java.util.List;

public class CustomSqlProviderMethodGenerator extends SqlProviderGenerator {

  public CustomSqlProviderMethodGenerator(String project) {
    super(project, false);
  }

  @Override
  public List<CompilationUnit> getCompilationUnits() {

    List<CompilationUnit> answer = super.getCompilationUnits();

    if (answer == null || answer.size() != 1) {
      return answer;
    }

    TopLevelClass topLevelClass = (TopLevelClass) answer.get(0);

    if (topLevelClass != null) {
      addInsertBatchMethod(topLevelClass);
    }

    return answer;
  }

  protected boolean addInsertBatchMethod(TopLevelClass topLevelClass) {

    AbstractJavaProviderMethodGenerator methodGenerator =
        new ProviderInsertBatchMethodGenerator(false);
    initializeAndExecuteGenerator(methodGenerator, topLevelClass);

    return true;
  }
}
